import numpy as nm
import scipy.sparse as sp
import datetime
import graph_Augmentation.sparseAugment as sa
import graph_Augmentation.randomAugment as ra
from logging import getLogger
from recbole.utils import set_color


# 整体方法说明
# 1、读入待增强的数据集
# 2、从交互数据中建立0-1交互邻接矩阵
# 3、进行全局增强加边,得到最终输出
def inter_random_augment(original_interactions, density_sample_ratio):
    # 程序运行计时
    start_time = datetime.datetime.now()

    # ==============第一部分：打印数据集的参数===================
    # 打印用户数、项目数、原矩阵中非零元素
    num_of_user = original_interactions.shape[0]
    num_of_item = original_interactions.shape[1]
    # print('num of user:', num_of_user)
    # print('num of item:', num_of_item)

    # print('original_interactions非零元素：', original_interactions.count_nonzero())
    logger = getLogger()
    logger.info(set_color('non-zero element in original_interactions', 'green')
                + f': {original_interactions.count_nonzero()}')

    # ==============第二部分：随机增强与资源矩阵计算===================
    # 生成随机增强矩阵，按数据集密度分别执行
    density = original_interactions.nnz / (num_of_user * num_of_item)
    # density_sample_ratio = 0.1
    logger.info(set_color('density', 'yellow') + f': {density}')
    logger.info(set_color('density_sample_ratio', 'green') + f': {density_sample_ratio}')

    if density > 0.01:
        random_augment_inter_matrix = ra.sample_random_matrix_augment(original_interactions,
                                                                      density_sample_ratio).tocoo()
    else:
        random_augment_inter_matrix = ra.large_random_matrix_augment(original_interactions,
                                                                     density_sample_ratio).tocoo()
    # print('random_augment非零元素：', random_augment_inter_matrix.count_nonzero())
    logger.info(set_color('non-zero element in random_augment', 'green')
                + f': {random_augment_inter_matrix.count_nonzero()}')

    # ==============第三部分：得到输出===================
    # 打印耗时
    # print(datetime.datetime.now(), ": 计算增强矩阵结束")
    logger.info(f'{datetime.datetime.now()}' + set_color(': calculate the enhancement matrix', 'green'))
    all_time = (datetime.datetime.now() - start_time).seconds
    # print("matrix_add_augment总耗时:", all_time)
    logger.info(set_color('all time of matrix_add_augment', 'green') + f': {all_time}')

    return random_augment_inter_matrix


# 整体方法说明
# 1、读入待增强的数据集
# 2、从交互数据中建立0-1交互邻接矩阵，注意从评分到矩阵元素转换有阈值要求；
# 3、在inter_graph基础上提取稀疏子图sparse_graph
# 4、进行全局增强加边
# 5、进行稀疏结点的针对性加边
# 6、把两个增强结果相加得到最终输出
def old_inter_graph_augment(original_interactions):
    # 程序运行计时
    start_time = datetime.datetime.now()

    # ==============第一部分：打印数据集的参数===================
    # 打印用户数、项目数、原矩阵中非零元素
    num_of_user = len(list(set(original_interactions.row)))
    num_of_item = len(list(set(original_interactions.col)))
    logger = getLogger()
    # print('num of user:', num_of_user)
    logger.info(set_color('num of user', 'green') + f': {num_of_user}')
    # print('num of item:', num_of_item)
    logger.info(set_color('num of item', 'green') + f': {num_of_item}')
    # print('original_interactions非零元素：', original_interactions.count_nonzero())
    logger.info(set_color('non-zero element in original_interactions', 'green') + f': {original_interactions.count_nonzero()}')
    # ==============第二部分：随机增强与资源矩阵计算===================
    # 生成随机增强矩阵，按稀疏度分流
    density = original_interactions.nnz / (num_of_user * num_of_item)
    if density > 0.1:
        random_augment_inter_matrix = ra.sample_random_matrix_augment(original_interactions)
    else:
        random_augment_inter_matrix = ra.large_random_matrix_augment(original_interactions)
    # print('random_augment非零元素：', random_augment_inter_matrix.count_nonzero())
    logger.info(set_color('non-zero element in random_augment', 'green')
                + f': {random_augment_inter_matrix.count_nonzero()}')

    # ==============第三部分：稀疏增强与资源矩阵计算===================
    # 生成稀疏增强子图
    sparse_inter_matrix = sa.sparse_matrix_augment(original_interactions)
    # print('sparse_inter非零元素：', sparse_inter_matrix.count_nonzero())
    logger.info(set_color('non-zero element in sparse_inter', 'green')
                + f': {sparse_inter_matrix.count_nonzero()}')

    # ==============第四部分：得到输出===================
    # 合并两个增强函数返回的增强矩阵
    aug_inter_matrix = random_augment_inter_matrix + sparse_inter_matrix
    # print(aug_inter_matrix)
    # 按阈值过滤,得到最终的增强矩阵
    gate = 1.5
    aug_filtered_inter_matrix = aug_filter(aug_inter_matrix, gate)
    # print('aug_filtered_inter_matrix非零元素：', aug_filtered_inter_matrix.count_nonzero())
    logger.info(set_color('non-zero element in aug_filtered_inter_matrix', 'green')
                + f': {aug_filtered_inter_matrix.count_nonzero()}')

    # 打印耗时
    # print(datetime.datetime.now(), ": 计算增强矩阵结束")
    logger.info(f'{datetime.datetime.now()}' + set_color(': calculate the enhancement matrix', 'green'))
    all_time = (datetime.datetime.now() - start_time).seconds
    # print("matrix_add_augment总耗时:", all_time)
    logger.info(set_color('all time of matrix_add_augment', 'green') + f': {all_time}')

    return aug_filtered_inter_matrix


# 按阈值对合并后的增强矩阵做过滤，返回coo_matrix
def aug_filter(sparse_matrix, gate):
    aug_tuple = sp.find(sparse_matrix >= gate)
    user_index = aug_tuple[0]
    item_index = aug_tuple[1]
    ratings = nm.ones_like(user_index, dtype=nm.float32)
    shape = sparse_matrix.shape
    aug_matrix = sp.coo_matrix((ratings, (user_index, item_index)), shape=shape)
    return aug_matrix


if __name__ == '__main__':
    old_inter_graph_augment('ml-1m')
